"""

    使用numpy 计算大型矩阵的平均值

"""

import numpy as np
import time
import asyncio
from concurrent.futures import ThreadPoolExecutor
from async_threading import async_timed
import functools


# def compute():
#     data_points = 4000000000
#
#     rows = 50
#
#     columns = int(data_points / rows)
#
#     matrix = np.arange(data_points).reshape(rows, columns)
#
#     s = time.time()
#
#     res = np.mean(matrix, axis=1)
#
#     print(time.time() - s)  # 59.16747498512268


def mean_for_row(arr, row):
    return np.mean(arr[row])


data_points = 4000000000

rows = 50

columns = int(data_points / rows)

matrix = np.arange(data_points).reshape(rows, columns)


@async_timed()  # 5.441076993942261s
async def main():
    loop = asyncio.get_running_loop()

    with ThreadPoolExecutor() as pool:
        tasks = []
        for i in range(rows):
            mean = functools.partial(mean_for_row, matrix, i)
            tasks.append(loop.run_in_executor(pool, mean))

        results = await asyncio.gather(*tasks)
        print(results)


asyncio.run(main())
